/**
 * @file libICA.c
 * 
 * Main FastICA functions.
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "libICA.h"



#define min(a,b) ((a<b) ? a : b)
#define abs(x) ((x)<0 ? -(x) : (x))

/**
 * Creates a matrix of given size.
 */
mat mat_create(int rows, int cols)
{
	mat M; int i;

	M = (mat) malloc(rows * sizeof(vect));
	if (M == NULL) {
		perror("Error allocating memory!");
		exit(-1);
	}
	for (i=0; i<rows; i++) {
		M[i] = (vect) malloc(cols * sizeof(scal));
		if (M[i] == NULL) {
			perror("Error allocating memory!");
			exit(-1);
		}
	}
	
	return M;
}

/**
 * Deletes matrix M.
 */
void mat_delete(mat M, int rows, int cols)
{
	int i;

	for (i=0; i<rows; i++)
		free(M[i]);
	free(M);
	M = NULL;
}

/**
 * Fills the whole matrix with 0's.
 */
void mat_zeroize(mat M, int rows, int cols)
{
	int i,j;

	for(i=0; i<rows; i++)
		for(j=0; j<cols; j++)
			M[i][j] = 0;
}

/**
 * Creates a vector of given size.
 */
vect vect_create(int n)
{
	vect v;

	v = (vect) malloc(n * sizeof(scal));
	if (v == NULL) {
		perror("Error allocating memory!");
		exit(-1);
	}	
	return v;
}

/**
 * Deletes vector v
 */
void vect_delete(vect v)
{
	free(v); v = NULL;
}

/**
 * Fills the whole vector with 0's.
 */
void vect_zeroize(vect v, int n)
{
	int i;

	for (i=0; i<n; i++)
		v[i] = 0;
}

/**
 * Applyies function fx() with parameter par on
 * every vector element.
 */
void vect_apply_fx(vect v, int n, scal (*fx)(double, scal), scal par)
{
	int i;
	
	for (i=0; i<n; i++)
		v[i] = (*fx)(v[i], par);
}

/**
 * Copyies mat M to Md.
 */
void mat_copy(mat M, int rows, int cols, mat Md)
{
	int i, j;

	for (i=0; i<rows; i++)
		for (j=0; j<cols; j++)
			Md[i][j] = M[i][j];
}

/**
 * Applyies function fx() with parameter par on
 * every matrix element.
 */
void mat_apply_fx(mat M, int rows, int cols, scal (*fx)(double, scal), scal par)
{
	int i,j;

	for (i=0; i<rows; i++)
		for (j=0; j<cols; j++)
			M[i][j] = (*fx)(M[i][j], par);
}

/**
 * Computes the mean of every row.
 */
void mat_mean_rows(mat M, int rows, int cols, vect v)
{
	int i,j;
	scal sum;

	for (i=0; i< rows; i++) {
		sum = 0;
		for (j=0; j<cols; j++)
			sum += M[i][j];
		v[i] = sum / cols;
	}
}

/**
 * Returns the maximal element on the diagonal
 * of the matrix M.
 */
scal mat_max_diag(mat M, int rows, int cols)
{
	int i; scal max;

	max = M[0][0];
	for (i=1; i<rows; i++)
		if (M[i][i] > max)
			max = M[i][i];
	return max;
}

/**
 * Creates a diagonal matrix from vector v.
 */
void mat_diag(vect v, int n, mat R)
{
	int i;

	mat_zeroize(R, n, n);
	for (i=0; i<n; i++)
		R[i][i] = v[i];
}

/**
 * Transponse matrix M.
 */
void mat_transpose(mat M, int rows, int cols, mat R)
{
	int i,j;

	for (i=0; i<rows; i++)
		for(j=0; j<cols; j++)
			R[j][i] = M[i][j]; 
}

/**
 * Compute inverse matrix. R = M^(-1).
 */
void mat_inverse(mat M, int dim, mat R)
{
	mat T; int i, j, k, maxrow;
	scal temp;

	// Create M|E	
	T = mat_create(dim, 2*dim);

	for (i=0; i<dim; i++)
		for (j=0; j<dim; j++)
			T[i][j] = M[i][j];
	for (i=0; i<dim; i++)
		for (j=dim; j<2*dim; j++)
			if (i == j - dim)
				T[i][j] = 1;
			else
				T[i][j] = 0;

	// Gauss-Jordan elimination
	for (i=0; i<dim; i++) {
		maxrow = i;
		for (j=i+1; j<dim; j++)
			if (abs(T[j][i]) > abs(T[maxrow][i]))
				maxrow = j;
		for (j=0; j<2*dim; j++) {
			temp = T[i][j];
			T[i][j] = T[maxrow][j];
			T[maxrow][j] = temp; 
		}
		if (abs(T[i][i]) <= SCAL_EPSILON) {
			perror("Inversion error: Singular matrix!");
			exit(-1);
		}
		for (j=i+1; j<dim; j++) {
			temp = T[j][i] / T[i][i];
			for (k=i; k<2*dim; k++)
				T[j][k] -= T[i][k] * temp;
		}
	}
	for (i=dim-1; i>=0; i--) {
		temp  = T[i][i];
		for (j=0; j<i; j++)
			for (k=2*dim-1; k>=i; k--)
				T[j][k] -=  T[i][k] * T[j][i] / temp;
		T[i][i] /= temp;
		for (j=dim; j<2*dim; j++) {
			T[i][j] /= temp;
			R[i][j-dim] = T[i][j];
		}
	}

	mat_delete(T, dim, 2*dim);
}

/**
 * Matrix subtraction. R = A - B.
 */
void mat_sub(mat A, mat B, int rows, int cols, mat R)
{
	int i,j;

	for (i=0; i<rows; i++)
		for (j=0; j<cols; j++)
			R[i][j] = A[i][j] - B[i][j];
}

/**
 * Matrix multiplication. R = A * B.
 */
void mat_mult(mat A, int rows_A, int cols_A, mat B, int rows_B, int cols_B, mat R)
{
	int i,j,k;

	mat_zeroize(R, rows_A, cols_B);
	for(i=0; i<rows_A; i++ )
		for(j=0; j<cols_B; j++)
			for(k=0; k<cols_A; k++)
				R[i][j] += A[i][k] * B[k][j];
}

/**
 * Centers mat M. (Subtracts the mean from every column)
 */
void mat_center(mat M, int rows, int cols, vect means)
{
	int i,j;

	vect_zeroize(means, cols);

	for (i=0; i<rows; i++)
		for(j=0; j<cols; j++)
			means[j] += M[i][j];		
	for (i=0; i<cols; i++)
		means[i] /= rows; 
	for (i=0; i<rows; i++)
		for(j=0; j<cols; j++)
			M[i][j] -= means[j];	
}

/**
 * De-center mat M. (Adds the mean to every column)
 */
void mat_decenter(mat M, int rows, int cols, vect means)
{
	int i,j;

	for(i=0; i<rows; i++)
		for(j=0; j<cols; j++)
			M[i][j] += means[j]; 
}



#define true  1
#define false 0


/**
 * Computes the singular value decomposition of A.
 * 
 * Give a matrix A, dimensions M by N, this routine computes its
 * singular value decomposition, A = U * W * transpose V. The matrix U
 * replaces A on output. The diagonal matrix of singular values, W, is
 * output as a vector W. The matrix V (not the transpose of V) is output
 * as V. M must be greater or equal to N. If it is smaller then A should
 * be filled up to square with zero rows.
 */
void svdcmp(mat A, int M, int N, vect W, mat V)
{
    /* Householder reduction to bidiagonal form. */
    int NM;
    scal C;
    scal F;
    scal G = 0.0;
    scal H;
    scal S;
    scal X;
    scal Y;
    scal Z;
    scal Scale = 0.0;
    scal ANorm = 0.0;
    scal tmp;
    int flag;
    int i;
    int its;
    int j;
    int jj;
    int k;
    int l;
	vect rv1;


	rv1 = vect_create(N);

    if( M < N ) {
        fprintf( stderr, "You must augment A with extra zero rows.\n" );
        return;
    }

    for( i = 0; i < N; ++i ) {
        l = i + 1;
        rv1[i] = Scale * G;
        G = 0.0;
        S = 0.0;
        Scale = 0.0;
        if( i < M ) {
            for( k = i; k < M; ++k ) {
                Scale = Scale + fabs( A[k][i] );
            }
            if( Scale != 0.0 ) {
                for( k = i; k < M; ++k ) {
                    A[k][i] = A[k][i] / Scale;
                    S = S + A[k][i] * A[k][i];
                }
                F = A[i][i];
                G = sqrt(S);
                if( F > 0.0 ) {
                    G = -G;
                }
                H = F * G - S;
                A[i][i] = F - G;
                if( i != (N-1) ) {
                    for( j = l; j < N; ++j ) {
                        S = 0.0;
                        for( k = i; k < M; ++k ) {
                            S = S + A[k][i] * A[k][j];
                        }
                        F = S / H;
                        for( k = i; k < M; ++k ) {
                            A[k][j] = A[k][j] + F * A[k][i];
                        }
                    }
                }
                for( k = i; k < M; ++k ) {
                    A[k][i] = Scale * A[k][i];
                }
            }
        }

        W[i] = Scale * G;
        G = 0.0;
        S = 0.0;
        Scale = 0.0;
        if( (i < M) && (i != (N-1)) ) {
            for( k = l; k < N; ++k ) {
                Scale = Scale + fabs( A[i][k] );
            }
            if( Scale != 0.0 ) {
                for( k = l; k < N; ++k ) {
                    A[i][k] = A[i][k] / Scale;
                    S = S + A[i][k] * A[i][k];
                }
                F = A[i][l];
                G = sqrt(S);
                if( F > 0.0 ) {
                    G = -G;
                }
                H = F * G - S;
                A[i][l] = F - G;
                for( k = l; k < N; ++k ) {
                    rv1[k] = A[i][k] / H;
                }
                if( i != (M-1) ) {
                    for( j = l; j < M; ++j ) {
                        S = 0.0;
                        for( k = l; k < N; ++k ) {
                            S = S + A[j][k] * A[i][k];
                        }
                        for( k = l; k < N; ++k ) {
                            A[j][k] = A[j][k] + S * rv1[k];
                        }
                    }
                }
                for( k = l; k < N; ++k ) {
                    A[i][k] = Scale * A[i][k];
                }
            }
        }
        tmp = fabs( W[i] ) + fabs( rv1[i] );
        if( tmp > ANorm )
            ANorm = tmp;
    }

    /* Accumulation of right-hand transformations. */
    for( i = N-1; i >= 0; --i ) {
        if( i < (N-1) ) {
            if( G != 0.0 ) {
                for( j = l; j < N; ++j ) {
                    V[j][i] = (A[i][j] / A[i][l]) / G;
                }
                for( j = l; j < N; ++j ) {
                    S = 0.0;
                    for( k = l; k < N; ++k ) {
                        S = S + A[i][k] * V[k][j];
                    }
                    for( k = l; k < N; ++k ) {
                        V[k][j] = V[k][j] + S * V[k][i];
                    }
                }
            }
            for( j = l; j < N; ++j ) {
                V[i][j] = 0.0;
                V[j][i] = 0.0;
            }
        }
        V[i][i] = 1.0;
        G = rv1[i];
        l = i;
    }

    /* Accumulation of left-hand transformations. */
    for( i = N-1; i >= 0; --i ) {
        l = i + 1;
        G = W[i];
        if( i < (N-1) ) {
            for( j = l; j < N; ++j ) {
                A[i][j] = 0.0;
            }
        }
        if( G != 0.0 ) {
            G = 1.0 / G;
            if( i != (N-1) ) {
                for( j = l; j < N; ++j ) {
                    S = 0.0;
                    for( k = l; k < M; ++k ) {
                        S = S + A[k][i] * A[k][j];
                    }
                    F = (S / A[i][i]) * G;
                    for( k = i; k < M; ++k ) {
                        A[k][j] = A[k][j] + F * A[k][i];
                    }
                }
            }
            for( j = i; j < M; ++j ) {
                A[j][i] = A[j][i] * G;
            }
        } else {
            for( j = i; j < M; ++j ) {
                A[j][i] = 0.0;
            }
        }
        A[i][i] = A[i][i] + 1.0;
    }

    /* Diagonalization of the bidiagonal form.
       Loop over singular values. */
    for( k = (N-1); k >= 0; --k ) {
        /* Loop over allowed iterations. */
        for( its = 1; its <= 30; ++its ) {
            /* Test for splitting.
               Note that rv1[0] is always zero. */
            flag = true;
            for( l = k; l >= 0; --l ) {
                NM = l - 1;
                if( (fabs(rv1[l]) + ANorm) == ANorm ) {
                    flag = false;
                    break;
                } else if( (fabs(W[NM]) + ANorm) == ANorm ) {
                    break;
                }
            }

            /* Cancellation of rv1[l], if l > 0; */
            if( flag ) {
                C = 0.0;
                S = 1.0;
                for( i = l; i <= k; ++i ) {
                    F = S * rv1[i];
                    if( (fabs(F) + ANorm) != ANorm ) {
                        G = W[i];
                        H = sqrt( F * F + G * G );
                        W[i] = H;
                        H = 1.0 / H;
                        C = ( G * H );
                        S = -( F * H );
                        for( j = 0; j < M; ++j ) {
                            Y = A[j][NM];
                            Z = A[j][i];
                            A[j][NM] = (Y * C) + (Z * S);
                            A[j][i] = -(Y * S) + (Z * C);
                        }
                    }
                }
            }
            Z = W[k];
            /* Convergence. */
            if( l == k ) {
                /* Singular value is made nonnegative. */
                if( Z < 0.0 ) {
                    W[k] = -Z;
                    for( j = 0; j < N; ++j ) {
                        V[j][k] = -V[j][k];
                    }
                }
                break;
            }

            if( its >= 30 ) {
                fprintf( stderr, "No convergence in 30 iterations.\n" );
                return;
            }

            X = W[l];
            NM = k - 1;
            Y = W[NM];
            G = rv1[NM];
            H = rv1[k];
            F = ((Y-Z)*(Y+Z) + (G-H)*(G+H)) / (2.0*H*Y);
            G = sqrt( F * F + 1.0 );
            tmp = G;
            if( F < 0.0 )
                tmp = -tmp;
            F = ((X-Z)*(X+Z) + H*((Y/(F+tmp))-H)) / X;

            /* Next QR transformation. */
            C = 1.0;
            S = 1.0;
            for( j = l; j <= NM; ++j ) {
                i = j + 1;
                G = rv1[i];
                Y = W[i];
                H = S * G;
                G = C * G;
                Z = sqrt( F * F + H * H );
                rv1[j] = Z;
                C = F / Z;
                S = H / Z;
                F = (X * C) + (G * S);
                G = -(X * S) + (G * C);
                H = Y * S;
                Y = Y * C;
                for( jj = 0; jj < N; ++jj ) {
                    X = V[jj][j];
                    Z = V[jj][i];
                    V[jj][j] = (X * C) + (Z * S);
                    V[jj][i] = -(X * S) + (Z * C);
                }
                Z = sqrt( F * F + H * H );
                W[j] = Z;

                /* Rotation can be arbitrary if Z = 0. */
                if( Z != 0.0 ) {
                    Z = 1.0 / Z;
                    C = F * Z;
                    S = H * Z;
                }
                F = (C * G) + (S * Y);
                X = -(S * G) + (C * Y);
                for( jj = 0; jj < M; ++jj ) {
                    Y = A[jj][j];
                    Z = A[jj][i];
                    A[jj][j] = (Y * C) + (Z * S);
                    A[jj][i] = -(Y * S) + (Z * C);
                }
            }
            rv1[l] = 0.0;
            rv1[k] = F;
            W[k] = X;
        }
    }
	vect_delete(rv1);
	
    return;
}

/*
 * Functions for matrix elements transformations
 * used in mat_apply_fx().
 */
static scal fx_inv(scal x, scal par)
{
	return (1/x);
}

static scal fx_inv_sqrt(scal x, scal par)
{
	return (1/sqrt(x));	
}

static scal fx_div_c(scal x, scal par)
{
	return (x/par);
}

static scal fx_rand(scal x, scal par)
{
	return (scal)rand()/RAND_MAX; 
}

static scal fx_tanh(scal x, scal par)
{
	return tanh(ALPHA * x);
}

static scal fx_1sub_sqr(scal x, scal par)
{
	return (ALPHA * (1-x*x));
}


/**
 * ICA function. Computes the W matrix from the
 * preprocessed data.
 */
static mat ICA_compute(mat X, int rows, int cols)
{
	mat TXp, GWX, W, Wd, W1, D, TU, TMP;
	vect d, lim;
	int i, it;

	// matrix creation
	TXp = mat_create(cols, rows);
	GWX = mat_create(rows, cols);
	W = mat_create(rows, rows);
	Wd = mat_create(rows, rows);
	D = mat_create(rows, rows);
	TMP = mat_create(rows, rows);
	TU = mat_create(rows, rows);
	W1 = mat_create(rows, rows);
	d = vect_create(rows);

	// W rand init
	mat_apply_fx(W, rows, rows, fx_rand, 0);

	// sW <- La.svd(W)
	mat_copy(W, rows, rows, Wd);
	svdcmp(Wd, rows, rows, d, D);

	// W <- sW$u %*% diag(1/sW$d) %*% t(sW$u) %*% W
	mat_transpose(Wd, rows, rows, TU);
	vect_apply_fx(d, rows, fx_inv, 0);
	mat_diag(d, rows, D);
	mat_mult(Wd, rows, rows, D, rows, rows, TMP);
	mat_mult(TMP, rows, rows, TU, rows, rows, D);
	mat_mult(D, rows, rows, W, rows, rows, Wd); // W = Wd

	// W1 <- W 
	mat_copy(Wd, rows, rows, W1);

	// lim <- rep(1000, maxit); it = 1
	lim = vect_create(MAX_ITERATIONS);
	for (i=0; i<MAX_ITERATIONS; i++)
		lim[i] = 1000;
	it = 0;


	// t(X)/p
	mat_transpose(X, rows, cols, TXp);
	mat_apply_fx(TXp, cols, rows, fx_div_c, cols);

	while (lim[it] > TOLERANCE && it < MAX_ITERATIONS) {
		// wx <- W %*% X
		mat_mult(Wd, rows, rows, X, rows, cols, GWX);

		// gwx <- tanh(alpha * wx)
		mat_apply_fx(GWX, rows, cols, fx_tanh, 0);
		
		// v1 <- gwx %*% t(X)/p
		mat_mult(GWX, rows, cols, TXp, cols, rows, TMP); // V1 = TMP
		
		// g.wx <- alpha * (1 - (gwx)^2)
		mat_apply_fx(GWX, rows, cols, fx_1sub_sqr, 0);

		// v2 <- diag(apply(g.wx, 1, FUN = mean)) %*% W
		mat_mean_rows(GWX, rows, cols, d);
		mat_diag(d, rows, D);
		mat_mult(D, rows, rows, Wd, rows, rows, TU); // V2 = TU

		// W1 <- v1 - v2
		mat_sub(TMP, TU, rows, rows, W1);
		
		// sW1 <- La.svd(W1)
		mat_copy(W1, rows, rows, W);
		svdcmp(W, rows, rows, d, D);

		// W1 <- sW1$u %*% diag(1/sW1$d) %*% t(sW1$u) %*% W1
		mat_transpose(W, rows, rows, TU);
		vect_apply_fx(d, rows, fx_inv, 0);
		mat_diag(d, rows, D);
		mat_mult(W, rows, rows, D, rows, rows, TMP);
		mat_mult(TMP, rows, rows, TU, rows, rows, D);
		mat_mult(D, rows, rows, W1, rows, rows, W); // W1 = W
		
		// lim[it + 1] <- max(Mod(Mod(diag(W1 %*% t(W))) - 1))
		mat_transpose(Wd, rows, rows, TU);
		mat_mult(W, rows, rows, TU, rows, rows, TMP);
		lim[it+1] = fabs(mat_max_diag(TMP, rows, rows) - 1);

		// W <- W1
		mat_copy(W, rows, rows, Wd);

		it++;
	}

	// clean up
	mat_delete(TXp, cols, rows);
	mat_delete(GWX, rows, cols);
	mat_delete(W, rows, rows);
	mat_delete(D, rows, rows);
	mat_delete(TMP, rows, rows);
	mat_delete(TU, rows, rows);
	mat_delete(W1, rows, rows);
	vect_delete(d);	

	return Wd;
}


/**
 * Main FastICA function. Centers and whitens the input
 * matrix, calls the ICA computation function ICA_compute()
 * and computes the output matrixes.
 */
void fastICA(mat X, int rows, int cols, int compc, mat K, mat W, mat A, mat S)
{
	mat XT, V, TU, D, X1, _A;
	vect scale, d;

	// matrix creation
	XT = mat_create(cols, rows);
	X1 = mat_create(compc, rows);
	V = mat_create(cols, cols);
	D = mat_create(cols, cols);
	TU = mat_create(cols, cols);
	scale = vect_create(cols);
	d = vect_create(cols);

	/*
	 * CENTERING
	 */
	mat_center(X, rows, cols, scale);


	/*
	 * WHITENING
	 */

	// X <- t(X); V <- X %*% t(X)/rows 
	mat_transpose(X, rows, cols, XT);
	mat_apply_fx(X, rows, cols, fx_div_c, rows);
	mat_mult(XT, cols, rows, X, rows, cols, V);
	
	// La.svd(V)
	svdcmp(V, cols, cols, d, D);  // V = s$u, d = s$d, D = s$v

	// D <- diag(c(1/sqrt(d))
	vect_apply_fx(d, cols, fx_inv_sqrt, 0);	
	mat_diag(d, cols, D);

	// K <- D %*% t(U)
	mat_transpose(V, cols, cols, TU);
	mat_mult(D, cols, cols, TU, cols, cols, V); // K = V 

	// X1 <- K %*% X
	mat_mult(V, compc, cols, XT, cols, rows, X1);

	/*
	 * FAST ICA
	 */
	_A = ICA_compute(X1, compc, rows);

	
	/*
	 * OUTPUT
	 */

	// X <- t(x)
	mat_transpose(XT, cols, rows, X);
	mat_decenter(X, rows, cols, scale);

	// K
	mat_transpose(V, compc, cols, K);

	// w <- a %*% K; S <- w %*% X
	mat_mult(_A, compc, compc, V, compc, cols, D);	
	mat_mult(D, compc, cols, XT, cols, rows, X1);
	// S
	mat_transpose(X1, compc, rows, S);

	// A <- t(w) %*% solve(w * t(w))
	mat_transpose(D, compc, compc, TU);
	mat_mult(D, compc, compc, TU, compc, compc, V);
	mat_inverse(V, compc, D);
	mat_mult(TU, compc, compc, D, compc, compc, V);
	// A
	mat_transpose(V, compc, compc, A);

	// W
	mat_transpose(_A, compc, compc, W);

	// cleanup
	mat_delete(XT, cols, rows);
	mat_delete(X1, compc, rows);
	mat_delete(V, cols, cols);
	mat_delete(D, cols, cols);
	mat_delete(TU,cols, cols);
	vect_delete(scale);
	vect_delete(d);
}

